import jax
import jax.numpy as jnp

import numpy as np

from min_ROC.filter_function import calcu_robust_infide_Vns

O, I, sx, sy, sz = (
	jnp.zeros((2, 2), dtype=complex),
	jnp.identity(2, dtype=complex),
	jnp.array([[0.0 + 0j, 1], [1, 0]]),
	jnp.array([[0.0, -1j], [1j, 0]]),
	jnp.array([[1.0 + 0j, 0], [0, -1]]),
)

def loss_fn(
		params,  # params
		func,
		dim,  # dimension of the system
		P,  # projection matrix

		ts,  # time-slots, N:=len(ts)-1
		ws,  # frequency slots

		H0,  # free part of Hamiltonian
		Hcs,  # control part of Hamiltonian, [Hc_1, Hc_2, ..., Hc_n], shape = (n, dim, dim)

		Vns,  # noise Hamiltonians, [Vn_1, Vn_2, ..., Vn_m],
		S_Vns_ws,  # PSD of noises
		alpha=1e-4,
		beta=1e-4,
		gamma=1.0,
):
	dts, dws = jnp.diff(ts), jnp.diff(ws)
	ucs_ts = func(params, ts)
	# psd_loss = calcu_robust_infide_Vns(ucs_ts, dim, P, ts, dts, ws, dws, H0, Hcs, Vns, S_Vns_ws)
	_, psd_loss = calcu_robust_infide_Vns(ucs_ts, dim, P, ts, dts, ws, dws, H0, Hcs, Vns, S_Vns_ws)
	dv_loss = (
			jnp.sum((jnp.diff(ucs_ts, axis=-1) / dts[None, :]) ** 2)
			+ jnp.sum((ucs_ts[:, 0] / dts[0]) ** 2)
			+ jnp.sum((ucs_ts[:, -1] / dts[-1]) ** 2)
	)
	max_amp_loss = jnp.max(ucs_ts ** 2)
	pulse_area = jnp.sum(0.5 * (ucs_ts[:, :-1] + ucs_ts[:, 1:]) * dts[None, :])
	gammas = jnp.ones(len(Vns)) * gamma
	loss = (
			(pulse_area - jnp.pi) ** 2
			+ alpha * dv_loss
			+ beta * max_amp_loss
			+ gammas @ psd_loss
	)
	return loss, (pulse_area, psd_loss, dv_loss)

import functools

def circ_pulse(params, ts, T, n_basis):
	w0_ts = 2.0 * jnp.pi / T * ts
	sin_ampl = params[1: n_basis + 1]
	cos_ampl = params[n_basis + 1:]
	ucs = jnp.sin(jnp.pi / T * ts) ** 2 * (
			params[0]
			+ sum(sin_ampl[i] * jnp.sin((i + 1.0) * w0_ts)
				  + cos_ampl[i] * jnp.cos((i + 1.0) * w0_ts)
				  for i in range(n_basis))
	)
	return jnp.asarray([ucs])

# ucs_ts = sin_pulse(params, ts[-1], ts, _n_basis=n_basis)[:-1]
N, Tf = 100, 100
w_max = 0.03 * (2 * jnp.pi)  # 30MHz
y, sigma = 1e-4 * w_max, 2e-1 * w_max  # 0.01MHz

w_span = jnp.logspace(-6, 0.2, 1000) * w_max
t_span = jnp.linspace(0, Tf, N + 1, endpoint=True)  # 0-500 ns

circ_pulse = functools.partial(circ_pulse, T=Tf, n_basis=3)

parameters = jnp.zeros(2 * 3 + 1)
parameters = parameters.at[0].set(2 * jnp.pi / Tf)

dynm = {
	"func"    : circ_pulse,
	"dim"     : 2,
	"ts"      : t_span,
	'H0'      : O,
	"Hcs"     : jnp.asarray([sx / 2.0]),
	'Vns'     : jnp.asarray([sz / 2.0]),
	"P"       : jnp.array([[1.0, 0], [0, 0]], dtype=complex),
	"ws"      : w_span,
	"S_Vns_ws": jnp.asarray([2 * sigma ** 2 * y / (y ** 2 + w_span ** 2)]),
}
import functools

loss_fn = functools.partial(loss_fn, **dynm)
loss_grad_fn = jax.jit(jax.value_and_grad(loss_fn, has_aux=True))

import optax
from tqdm import trange

pbar = trange(1000, desc="Iterations", leave=True)
lr = optax.exponential_decay(1e-3, 1, 0.999, 500)
solver = optax.adam(learning_rate=lr)

losses = []
opt_state = solver.init(parameters)
for _ in pbar:
	(loss, (pulse_area, psd_loss, dv_loss)), grad = loss_grad_fn(parameters)

	updates, opt_state = solver.update(grad, opt_state, parameters)
	parameters = optax.apply_updates(parameters, updates)

	pbar.set_postfix(
			{"loss": loss, "pulse_area": pulse_area, "psd_loss": psd_loss, "dv_loss": dv_loss}
	)
	losses.append(psd_loss[0])
import matplotlib.pyplot as plt

plt.semilogy(np.arange(len(losses)), losses)
plt.show()
# plt.step(t_span[:-1], params, where='post')
plt.plot(t_span, circ_pulse(parameters, t_span)[0])
plt.show()
print(parameters)
